#include <cstdio>
#include <cstdint>
#include <cassert>
#include <cstring>
#include <vector>
#include <iterator>
#include <fstream>

#define VIRUS_SIZE 41472
#define INFECTED_SIZE 1000
#define KEY_OFFSET 1234

static inline void imul32(uint32_t * dest, uint32_t d, uint32_t s)
{
	int64_t res = (int64_t)(int32_t)d * (int32_t)s;
	*dest = (uint32_t)res;
}

static inline void mul32(uint32_t * dest, uint32_t d, uint32_t s)
{
	uint64_t res = (uint64_t)d * s;
	*dest = ((uint32_t)(res >> 32));
}

void bytes_decode(char * data, int count, uint32_t key)
{
	//printf("key: 0x%08x\n", key);
	uint32_t res = 0;
	do {
		imul32(&key, key, 0x8088405);
		key++;

		mul32(&res, key, 0xff);
		*data++ ^= res;
	} while(--count);
}

void neshta_fix(const char * exe_path)
{
	std::ifstream ifs(exe_path, std::ios::binary);
	if (!ifs)
		throw "Error opening file ";
	std::vector<char> infected_exe = std::vector<char>(std::istreambuf_iterator<char>(ifs), {});
	ifs.close();

	if(infected_exe.size() <= 41472 || *((uint32_t *)(infected_exe.data() + 0x1a89)) != 0x8088405)
		return;

	char * exe_data = infected_exe.data();
	unsigned long exe_size = infected_exe.size();
	char * end = exe_data + exe_size - VIRUS_SIZE;
	printf("Repair[%s]\n", exe_path);
	bytes_decode(end, INFECTED_SIZE, *((uint32_t*)(end + KEY_OFFSET)));
	memcpy(exe_data, end, VIRUS_SIZE);

	FILE * out_fp = fopen(exe_path, "wb");
	assert(out_fp != NULL);
	fwrite(exe_data, sizeof(char), (exe_size - VIRUS_SIZE), out_fp);
	fclose(out_fp);
}

int main(int argc, const char ** argv)
{
	if(argc != 2)
	{
		printf("Usage: %s [infected_exe]\n", argv[0]);
		return 1;
	}

	neshta_fix(argv[1]);

	return 0;
}
